"""
WebSocket管理模块

提供实时数据推送功能
"""

import json
import asyncio
import threading
from datetime import datetime
from typing import Dict, Any, Set, Optional
import websockets
from websockets.server import WebSocketServerProtocol

from ..core import unified_logger
from ..monitoring import unified_monitor


class WebSocketManager:
    """WebSocket管理器"""
    
    def __init__(self, monitor=None, host: str = 'localhost', port: int = 8765):
        """
        初始化WebSocket管理器
        
        Args:
            monitor: 监控实例
            host: 服务器主机
            port: 服务器端口
        """
        self.monitor = monitor or unified_monitor
        self.host = host
        self.port = port
        self.logger = unified_logger.get_logger('websocket_manager')
        
        # 连接管理
        self.clients: Set[WebSocketServerProtocol] = set()
        self.subscriptions: Dict[WebSocketServerProtocol, Set[str]] = {}
        
        # 服务器状态
        self.server = None
        self.running = False
        self.thread = None
        
        # 数据推送
        self.push_interval = 5  # 推送间隔（秒）
        self.push_task = None
        
        self.logger.info(f"WebSocket管理器初始化完成: {host}:{port}")
    
    def start(self):
        """启动WebSocket服务器"""
        if self.running:
            self.logger.warning("WebSocket服务器已在运行")
            return
        
        self.running = True
        self.thread = threading.Thread(target=self._run_server, daemon=True)
        self.thread.start()
        
        self.logger.info("WebSocket服务器启动中...")
    
    def stop(self):
        """停止WebSocket服务器"""
        if not self.running:
            return
        
        self.running = False
        
        # 停止推送任务
        if self.push_task:
            self.push_task.cancel()
        
        # 关闭所有连接
        if self.clients:
            asyncio.create_task(self._close_all_connections())
        
        # 停止服务器
        if self.server:
            self.server.close()
        
        self.logger.info("WebSocket服务器已停止")
    
    def _run_server(self):
        """运行WebSocket服务器"""
        try:
            # 创建事件循环
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            # 启动服务器
            start_server = websockets.serve(
                self._handle_client,
                self.host,
                self.port
            )
            
            self.server = loop.run_until_complete(start_server)
            
            # 启动数据推送任务
            self.push_task = loop.create_task(self._push_data_loop())
            
            self.logger.info(f"WebSocket服务器已启动: ws://{self.host}:{self.port}")
            
            # 运行事件循环
            loop.run_forever()
            
        except Exception as e:
            self.logger.error(f"WebSocket服务器运行异常: {e}")
        finally:
            self.running = False
    
    async def _handle_client(self, websocket: WebSocketServerProtocol, path: str):
        """处理客户端连接"""
        client_addr = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
        self.logger.info(f"新的WebSocket连接: {client_addr}")
        
        # 添加客户端
        self.clients.add(websocket)
        self.subscriptions[websocket] = set()
        
        try:
            # 发送欢迎消息
            await self._send_message(websocket, {
                'type': 'welcome',
                'message': 'Connected to XCQuant WebSocket',
                'timestamp': datetime.now().isoformat()
            })
            
            # 处理消息
            async for message in websocket:
                await self._handle_message(websocket, message)
                
        except websockets.exceptions.ConnectionClosed:
            self.logger.info(f"WebSocket连接关闭: {client_addr}")
        except Exception as e:
            self.logger.error(f"处理WebSocket连接异常: {e}")
        finally:
            # 清理客户端
            self.clients.discard(websocket)
            self.subscriptions.pop(websocket, None)
    
    async def _handle_message(self, websocket: WebSocketServerProtocol, message: str):
        """处理客户端消息"""
        try:
            data = json.loads(message)
            msg_type = data.get('type')
            
            if msg_type == 'subscribe':
                # 订阅数据类型
                channels = data.get('channels', [])
                for channel in channels:
                    self.subscriptions[websocket].add(channel)
                
                await self._send_message(websocket, {
                    'type': 'subscribed',
                    'channels': list(self.subscriptions[websocket]),
                    'timestamp': datetime.now().isoformat()
                })
                
            elif msg_type == 'unsubscribe':
                # 取消订阅
                channels = data.get('channels', [])
                for channel in channels:
                    self.subscriptions[websocket].discard(channel)
                
                await self._send_message(websocket, {
                    'type': 'unsubscribed',
                    'channels': list(self.subscriptions[websocket]),
                    'timestamp': datetime.now().isoformat()
                })
                
            elif msg_type == 'ping':
                # 心跳检测
                await self._send_message(websocket, {
                    'type': 'pong',
                    'timestamp': datetime.now().isoformat()
                })
                
            elif msg_type == 'get_data':
                # 获取即时数据
                channel = data.get('channel')
                if channel:
                    channel_data = await self._get_channel_data(channel)
                    await self._send_message(websocket, {
                        'type': 'data',
                        'channel': channel,
                        'data': channel_data,
                        'timestamp': datetime.now().isoformat()
                    })
                
        except json.JSONDecodeError:
            await self._send_message(websocket, {
                'type': 'error',
                'message': 'Invalid JSON format',
                'timestamp': datetime.now().isoformat()
            })
        except Exception as e:
            self.logger.error(f"处理消息异常: {e}")
            await self._send_message(websocket, {
                'type': 'error',
                'message': str(e),
                'timestamp': datetime.now().isoformat()
            })
    
    async def _send_message(self, websocket: WebSocketServerProtocol, message: Dict[str, Any]):
        """发送消息到客户端"""
        try:
            await websocket.send(json.dumps(message, ensure_ascii=False))
        except websockets.exceptions.ConnectionClosed:
            pass
        except Exception as e:
            self.logger.error(f"发送消息失败: {e}")
    
    async def _broadcast_message(self, message: Dict[str, Any], channel: str = None):
        """广播消息到所有订阅的客户端"""
        if not self.clients:
            return
        
        # 过滤订阅了指定频道的客户端
        target_clients = []
        for client in self.clients.copy():
            if channel is None or channel in self.subscriptions.get(client, set()):
                target_clients.append(client)
        
        if not target_clients:
            return
        
        # 并发发送消息
        tasks = []
        for client in target_clients:
            tasks.append(self._send_message(client, message))
        
        await asyncio.gather(*tasks, return_exceptions=True)
    
    async def _push_data_loop(self):
        """数据推送循环"""
        while self.running:
            try:
                await asyncio.sleep(self.push_interval)
                
                if not self.clients:
                    continue
                
                # 收集需要推送的数据类型
                all_channels = set()
                for subscriptions in self.subscriptions.values():
                    all_channels.update(subscriptions)
                
                # 推送各类数据
                for channel in all_channels:
                    try:
                        data = await self._get_channel_data(channel)
                        if data:
                            await self._broadcast_message({
                                'type': 'data',
                                'channel': channel,
                                'data': data,
                                'timestamp': datetime.now().isoformat()
                            }, channel)
                    except Exception as e:
                        self.logger.error(f"推送{channel}数据失败: {e}")
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"数据推送循环异常: {e}")
    
    async def _get_channel_data(self, channel: str) -> Optional[Dict[str, Any]]:
        """获取频道数据"""
        try:
            if channel == 'dashboard':
                # 监控面板数据
                return self.dashboard.get_dashboard_data()
            
            elif channel == 'system':
                # 系统指标
                if self.dashboard.system_monitor:
                    current_metrics = self.dashboard.system_monitor.get_current_metrics()
                    if current_metrics:
                        return {
                            'cpu_percent': current_metrics.cpu_percent,
                            'memory_percent': current_metrics.memory_percent,
                            'disk_percent': current_metrics.disk_percent,
                            'process_count': current_metrics.process_count,
                            'timestamp': current_metrics.timestamp.isoformat()
                        }
            
            elif channel == 'strategies':
                # 策略状态
                if self.dashboard.strategy_monitor:
                    strategies = self.dashboard.strategy_monitor.get_all_strategies_status()
                    performance = self.dashboard.strategy_monitor.get_performance_summary()
                    return {
                        'strategies': strategies,
                        'performance': performance
                    }
            
            elif channel == 'alerts':
                # 告警信息
                if self.dashboard.alert_manager:
                    active_alerts = self.dashboard.alert_manager.get_active_alerts()
                    stats = self.dashboard.alert_manager.get_stats()
                    return {
                        'active_alerts': [
                            {
                                'id': alert.id,
                                'title': alert.title,
                                'message': alert.message,
                                'level': alert.level.value,
                                'status': alert.status.value,
                                'created_at': alert.created_at.isoformat(),
                                'tags': alert.tags
                            }
                            for alert in active_alerts
                        ],
                        'stats': stats
                    }
            
            elif channel == 'health':
                # 系统健康状态
                return self.dashboard.get_system_health()
            
            return None
            
        except Exception as e:
            self.logger.error(f"获取{channel}数据失败: {e}")
            return None
    
    async def _close_all_connections(self):
        """关闭所有连接"""
        if not self.clients:
            return
        
        # 发送关闭通知
        close_message = {
            'type': 'server_shutdown',
            'message': 'Server is shutting down',
            'timestamp': datetime.now().isoformat()
        }
        
        tasks = []
        for client in self.clients.copy():
            tasks.append(self._send_message(client, close_message))
        
        await asyncio.gather(*tasks, return_exceptions=True)
        
        # 关闭连接
        tasks = []
        for client in self.clients.copy():
            tasks.append(client.close())
        
        await asyncio.gather(*tasks, return_exceptions=True)
        
        self.clients.clear()
        self.subscriptions.clear()
    
    def get_status(self) -> Dict[str, Any]:
        """获取WebSocket服务器状态"""
        return {
            'running': self.running,
            'host': self.host,
            'port': self.port,
            'connected_clients': len(self.clients),
            'total_subscriptions': sum(len(subs) for subs in self.subscriptions.values()),
            'push_interval': self.push_interval
        }
    
    async def send_notification(self, title: str, message: str, level: str = 'info'):
        """发送通知到所有客户端"""
        notification = {
            'type': 'notification',
            'title': title,
            'message': message,
            'level': level,
            'timestamp': datetime.now().isoformat()
        }
        
        await self._broadcast_message(notification)